from PyQt6.QtWidgets import QInputDialog
import os
from openpyxl import Workbook, load_workbook
from openpyxl.styles import Alignment
from openpyxl.utils import get_column_letter
from datetime import datetime  # 导入 datetime 模块

class SaveResults:
    def __init__(self, parent):
        self.parent = parent
        self.output_folder = parent.output_folder
        self.excel_file_name = parent.excel_file_name
        self.current_results = parent.current_results
        # self.FIELD_ORDER = ['基本信息', '教育背景', '实习经历', '技能证书']  # 添加 self. 前缀
        self.FIELD_ORDER = []  # 初始化为空列表，动态生成

    # 新增方法：动态生成字段顺序
    def generate_field_order(self, structured_list):
        field_order = []
        for data in structured_list:
            for main_field in data.keys():
                if main_field not in field_order:
                    field_order.append(main_field)
        return field_order

    def update_output_folder(self, new_folder):
        self.output_folder = new_folder

    def save_results(self):
        format, ok = QInputDialog.getItem(
            self.parent, "选择保存格式",
            "请选择保存格式：",
            ["Excel (.xlsx)", "Text (.txt)"],
            0, False
        )
        if not ok:
            return

        if format == "Excel (.xlsx)":
            self.save_to_excel()
        elif format == "Text (.txt)":
            self.save_to_text()

    def save_to_excel(self):
        if not self.current_results:
            self.parent.show_warning("保存失败", "没有可保存的识别结果")
            return

        try:
            excel_path = os.path.join(self.output_folder, self.excel_file_name)

            # 解析结构化数据
            structured_data = self.parse_structured_data()  # 新增解析方法

            # 写入Excel
            self.write_structured_excel(structured_data, excel_path)  # 重构写入方法

            self.parent.show_info("保存成功", f"文件已保存到：\n{excel_path}")

        except PermissionError:
            self.parent.show_error("保存失败", "文件被其他程序占用，请关闭Excel后重试")
        except Exception as e:
            self.parent.show_error("保存失败", f"写入Excel文件时出错：\n{str(e)}")

    # 重构Excel写入方法 --------------------------------------------
    def write_structured_excel(self, structured_list, excel_path):
        if os.path.exists(excel_path):
            wb = load_workbook(excel_path)
            ws = wb.active
            is_new_file = False
        else:
            wb = Workbook()
            ws = wb.active
            ws.append([])  # 初始空行
            is_new_file = True

        # 写入标题（仅首次创建时）
        if is_new_file:
            current_col = 1
            for main_field in self.FIELD_ORDER:
                # 获取所有图片的第一个结构化数据作为标题参考（假设字段一致）
                sample_data = structured_list[0] if structured_list else {}
                sub_fields = sample_data.get(main_field, {})
                end_col = current_col + len(sub_fields) - 1

                # 合并主字段标题
                ws.merge_cells(start_row=1, start_column=current_col, end_row=1, end_column=end_col)
                title_cell = ws.cell(row=1, column=current_col, value=main_field)
                title_cell.alignment = Alignment(horizontal='center', vertical='center')

                # 写入子字段标题
                for idx, sub_field in enumerate(sub_fields.keys(), start=current_col):
                    header_cell = ws.cell(row=2, column=idx, value=sub_field)
                    header_cell.alignment = Alignment(horizontal='center', vertical='center')

                current_col = end_col + 2  # 字段间留空列
            start_row = 3
        else:
            start_row = ws.max_row + 1  # 追加数据时从最后一行开始

        # 按图片逐个写入数据
        for structured_data in structured_list:
            current_col = 1
            for main_field in self.FIELD_ORDER:
                sub_fields = structured_data.get(main_field, {})
                end_col = current_col + len(sub_fields) - 1

                # 填充数据到对应列
                for idx, value in enumerate(sub_fields.values(), start=current_col):
                    cell = ws.cell(row=start_row, column=idx, value=value)
                    cell.alignment = Alignment(wrap_text=True, vertical='center')

                current_col = end_col + 2  # 移动到下一个主字段起始列

            start_row += 1  # 换行处理下一个图片数据

        # 动态调整列宽（保持原有逻辑）
        for col_num in range(1, ws.max_column + 1):
            max_len = 0
            for row in ws.iter_rows(min_row=2, max_row=ws.max_row, min_col=col_num, max_col=col_num):
                cell = row[0]
                if cell.value:
                    lines = str(cell.value).split('\n')
                    for line in lines:
                        ch_width = sum(2 if '\u4e00' <= c <= '\u9fff' else 1 for c in line)
                        max_len = max(max_len, ch_width)

            adjusted_width = max_len * 1.1 + 1
            ws.column_dimensions[get_column_letter(col_num)].width = adjusted_width

        wb.save(excel_path)

    # 新增结构化解析方法 --------#按图片逐个解析，返回结构化数据列表------------
    def parse_structured_data(self):
        structured_list = []
        for path, text_lines in self.current_results.items():
            raw_lines = [line.strip() for line in text_lines if line.strip()]
            structured = self.enhanced_parse(raw_lines)
            structured_list.append(structured)
        # 动态生成 FIELD_ORDER
        self.FIELD_ORDER = self.generate_field_order(structured_list)
        return structured_list    # 返回列表，每个元素对应一个图片的结构化数据


    def enhanced_parse(self, lines):
        structured = {}
        current_main_field = None
        pending_field = None
        buffer = []

        for line in lines:
            line = line.strip()
            if not line:
                continue

            # 判断是否为大字段名：当前行文本数量较少（3~5个字符）且没有标点符号
            if len(line) <= 5 and not any(char in line for char in ['。', '，', '：', ':', '；', ';']):
                if current_main_field and pending_field and buffer:
                    structured[current_main_field][pending_field] = '\n'.join(buffer)
                    buffer = []
                current_main_field = line  # 将当前行作为大字段名
                structured[current_main_field] = {}  # 初始化大字段
                pending_field = None
                continue

            # 处理子字段
            if '：' in line or ':' in line:
                sep = '：' if '：' in line else ':'
                key_part, _, value_part = line.partition(sep)
                key = key_part.strip()
                value = value_part.strip()

                if pending_field and buffer:
                    structured[current_main_field][pending_field] = '\n'.join(buffer)
                    buffer = []

                if value:
                    structured[current_main_field][key] = value
                    pending_field = None
                else:
                    pending_field = key
            else:
                # 处理多行文本（如“主要职责”和“项目成果”）
                if pending_field:
                    buffer.append(line)
                elif current_main_field and structured.get(current_main_field):
                    last_key = list(structured[current_main_field].keys())[-1]
                    # 如果最后一个字段已经有值，则追加到该字段的值中
                    if structured[current_main_field][last_key]:
                        structured[current_main_field][last_key] += '\n' + line
                    else:
                        structured[current_main_field][last_key] = line

        # 处理最后一个 pending_field 和 buffer
        if pending_field and buffer:
            structured[current_main_field][pending_field] = '\n'.join(buffer)

        return structured

    def save_to_text(self):
        if not self.current_results:
            self.parent.show_warning("保存失败", "没有可保存的识别结果")
            return

        try:
            # 获取当前时间并格式化为字符串
            current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            # 生成文件路径
            base_name = self.excel_file_name.replace(".xlsx", "")
            text_path = os.path.join(self.output_folder, f"{base_name}.txt")

            with open(text_path, "a", encoding="utf-8") as file:
                # 在文件开头添加时间戳
                file.write(f"=== 保存时间：{current_time} ===\n\n")
                for path, text in self.current_results.items():
                    file.write(f"=== 图片路径：{path} ===\n")
                    file.write(f"{text}\n\n")

            self.parent.show_info("保存成功", f"文件已保存到：\n{text_path}")

        except Exception as e:
            self.parent.show_error("保存失败", f"写入文本文件时出错：\n{str(e)}")